package org.example.doris.arrow.adbc;

import org.apache.arrow.adbc.core.*;
import org.apache.arrow.adbc.driver.flightsql.FlightSqlDriver;
import org.apache.arrow.flight.Location;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowReader;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ADBCDriverDemo {
    public static void main(String[] args) throws Exception {
        // 1. new driver
        final BufferAllocator allocator = new RootAllocator();
        FlightSqlDriver driver = new FlightSqlDriver(allocator);
        Map<String, Object> parameters = new HashMap<>();
        AdbcDriver.PARAM_URI.set(parameters, Location.forGrpcInsecure("192.168.0.227", 9090).getUri().toString());
        AdbcDriver.PARAM_USERNAME.set(parameters, "root");
        AdbcDriver.PARAM_PASSWORD.set(parameters, "123456");
        AdbcDatabase adbcDatabase = driver.open(parameters);

        // 2. new connection
        AdbcConnection connection = adbcDatabase.connect();
        AdbcStatement stmt = connection.createStatement();

        // 3. execute query
        // stmt.setSqlQuery("select * from information_schema.tables;");
        stmt.setSqlQuery("select * from testdb.table_hash;");
        AdbcStatement.QueryResult queryResult = stmt.executeQuery();
        ArrowReader reader = queryResult.getReader();

        // 4. load result
        List<String> result = new ArrayList<>();
        while (reader.loadNextBatch()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            String tsvString = root.contentToTSVString();
            result.add(tsvString);
            System.out.println(tsvString);
        }
        System.out.printf("batchs %d\n", result.size());

        // 5. close
        reader.close();
        queryResult.close();
        stmt.close();
        connection.close();
    }
}
